package com.chinaums.sftp.source;

import com.chinaums.sftp.internal.connection.SftpHandlerImpl;
import com.chinaums.sftp.internal.constants.ConstantValue;
import com.chinaums.sftp.internal.options.SftpReadOptions;
import com.chinaums.sftp.tablesource.SftpSeqBufferedReader;
import com.chinaums.sftp.internal.connection.SftpHandler;
import com.chinaums.sftp.utils.FileUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.table.data.RowData;
import org.apache.flink.util.UserCodeClassLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.Statement;
import java.util.*;

import static com.chinaums.sftp.internal.constants.ConstantValue.URL;
import static com.chinaums.sftp.internal.constants.ConstantValue.NAME;
import static com.chinaums.sftp.internal.constants.ConstantValue.PASSWORD;
import static com.chinaums.sftp.internal.constants.ConstantValue.DRIVER;
import static com.chinaums.sftp.internal.constants.ConstantValue.TABLE;
import static com.chinaums.sftp.internal.constants.ConstantValue.isMetric;


/**
 * stream source function.
 */
@Deprecated
@Internal
public class SftpSourceFunction extends RichSourceFunction<RowData> implements ResultTypeQueryable<RowData>, CheckpointedFunction {
    private static final long serialVersionUID = 1L;
    private static final Logger LOG = LoggerFactory.getLogger(SftpSourceFunction.class);

    private final SftpReadOptions sftpReadOptions;
    private final DeserializationSchema<RowData> deserialization;
    private final TypeInformation<RowData> rowDataTypeInfo;
    private transient boolean isErr = false;

    private transient SftpHandler ftpHandler;
    private transient SftpSeqBufferedReader br;
    private volatile boolean isRunning = true;
    private transient ListState<Tuple2<String, Long>> fileReadOffset;
    private transient Map<String, Long> fileReadIndex;

    private transient Long exitIndex = -1L;

    private transient Counter counter;

    private transient HashMap<String, Counter> toMysql;

    // 多并行度问题 一个线程一个文件 .... 其他线程关闭即可.
    private transient boolean tag = false;

    private transient Connection conn;

    private transient Statement stmt;


    public SftpSourceFunction(
            SftpReadOptions sftpReadOptions,
            DeserializationSchema<RowData> deserialization,
            TypeInformation<RowData> rowDataTypeInfo) {
        this.sftpReadOptions = sftpReadOptions;
        this.deserialization = deserialization;
        this.rowDataTypeInfo = rowDataTypeInfo;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        deserialization.open(new DeserializationSchema.InitializationContext() {
            @Override
            public MetricGroup getMetricGroup() {
                return getRuntimeContext().getMetricGroup();
            }

            @Override
            public UserCodeClassLoader getUserCodeClassLoader() {
                return null;
            }
        });
        ftpHandler = new SftpHandlerImpl();
        ftpHandler.loginFtpServer(sftpReadOptions);

        if (isMetric() && sftpReadOptions.isMetric()) {
            Class.forName(DRIVER);
            conn = DriverManager.getConnection(URL, NAME, PASSWORD);
            stmt = conn.createStatement();
            this.counter = getRuntimeContext().getMetricGroup().counter("SFTPCounter");
            toMysql = new HashMap<>();
        }

    }

    @Override
    public TypeInformation<RowData> getProducedType() {
        return rowDataTypeInfo;
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        if (isRunning) {
            fileReadOffset.clear();
            if (br != null) {
                fileReadIndex = br.getFileReadIndex();
            }
            fileReadIndex.keySet().forEach(key -> {
                Tuple2<String, Long> t = new Tuple2<>();
                t.f0 = key;
                t.f1 = fileReadIndex.get(key);
                try {
                    fileReadOffset.add(t);
                    LOG.info("store checkPointed file index msg - " + t.f0 + ":" + t.f1);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            });
        }
    }

    @Override
    public void initializeState(FunctionInitializationContext ctx) throws Exception {
        ListStateDescriptor<Tuple2<String, Long>> stateDescriptor = new ListStateDescriptor<>(
                "file read index for single list",
                TypeInformation.of(new TypeHint<Tuple2<String, Long>>() {
                }));
        fileReadOffset = ctx.getOperatorStateStore().getListState(stateDescriptor);
        fileReadIndex = new HashMap<>();
        if (ctx.isRestored()) {
            for (Tuple2<String, Long> t : fileReadOffset.get()) {
                fileReadIndex.put(t.f0, t.f1);
                LOG.info(String.format("restore checkPointed file name - %s size - %s", t.f0, t.f1));
            }
        }
    }

    @Override
    public void run(SourceContext<RowData> ctx) throws Exception {
        LOG.error("########################"+isRunning);
        Loop:
        while (isRunning) {
            List<Tuple2<String, Long>> subtaskFiles = new ArrayList<>();

            List<String> files = new ArrayList<>();
            if (!isRunning) {
                break;
            }
            FileUtil.addFiles(
                    sftpReadOptions.getPath(),
                    ftpHandler,
                    files);
            String fileName ="";
            // 遍历文件使一个线程一个文件 ....
            for (String file : files) {
                if (file.contains(ConstantValue.FILE_REGEX_SPLIT)) {
                    String[] fileInfo = file.split(ConstantValue.FILE_REGEX_SPLIT);
                    tag = true;
                    FileUtil.createFileInfo(subtaskFiles, fileInfo, fileReadIndex);
                }
            }
            // 退出条件 1.多余线程,2.没有运行的线程, 3.已经完成的线程
            if (!tag || !isRunning || SftpSeqBufferedReader.isTail || subtaskFiles.size() == 0 || Objects.equals(exitIndex, subtaskFiles.get(0).f1)) {
                break;
            } else {
                exitIndex = subtaskFiles.get(0).f1;
            }


            // 实例 SftpSeqBufferedReader
            br = new SftpSeqBufferedReader(
                    ftpHandler,
                    subtaskFiles.stream().map(s -> s.f0).iterator(),
                    sftpReadOptions
            );

            // 从指定位置读取数据
            if (fileReadIndex != null && fileReadIndex.size() > 0) {
                br.setFileReadIndex(fileReadIndex);
            }
            // 首行跳过...
            if (sftpReadOptions.isFirstLineHeader()) {
                br.setFromLine(1);
            } else {
                br.setFromLine(0);
            }
            // 数据格式...
            br.setFileEncoding(sftpReadOptions.getEncoding());
            try {
                while (true) {
                    if (!isRunning) {
                        break Loop;
                    }
                    String line;
                    synchronized (ctx.getCheckpointLock()) {
                        if ((line = br.readLine(false)) == null) {
                            break;
                        }
                    }
                    // 自定义序列化并转化成RowData数据
                    RowData row = deserialization.deserialize(line.getBytes(StandardCharsets.UTF_8));
                    //  监控数据
                    if (row != null) {
                        if (isMetric() && sftpReadOptions.isMetric()) {
                            if (br.getFileName().equals(fileName)  || StringUtils.isEmpty(fileName) ){
                                counter.inc();
                                toMysql.put(br.getFileName(), counter);
                            }else {
                                this.counter = getRuntimeContext().getMetricGroup().counter("SFTPCounter");
                                counter.inc();
                                toMysql.put(br.getFileName(), counter);
                            }
                        }
                        fileName = br.getFileName();
                        ctx.collect(row);
                    }
                }
            } catch (Exception e) {
                isErr = true;
                throw new RuntimeException(e);
            }
            if (br != null) {
                fileReadIndex = br.getFileReadIndex();
                br.close();
                br = null;
            }
            Thread.sleep(sftpReadOptions.getStreamInterval());
        }
    }

    @Override
    public void cancel() {
        try {
            isRunning = false;
            if (br != null) {
                br.close();
            }
            if (ftpHandler != null) {
                ftpHandler.logoutFtpServer();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     *  监控数据写入到mysql
     */
    @Override
    public void close() throws Exception {
        //添加try/catch 是 防止mysql 连接失败 影响整体任务的运行。
        try {
            if (isMetric() && sftpReadOptions.isMetric() && tag && (!isErr)) {
                Class.forName(DRIVER);
                conn = DriverManager.getConnection(URL, NAME, PASSWORD);
                stmt = conn.createStatement();
                String status = isErr ? "err" : "ok";
                for (Map.Entry<String, Counter> counterEntry : toMysql.entrySet()) {
                    String sql = "INSERT INTO " + TABLE + " (status, file_name, thread_num, count) VALUES(  '" + status + "' , '" + counterEntry.getKey() + "', " + this.getRuntimeContext().getIndexOfThisSubtask() + ", " + counterEntry.getValue().getCount() + ")";
                    stmt.execute(sql);
                }
                stmt.close();
                conn.close();
            }
        } catch (Exception ignored) {
            LOG.error("mysql监控入库异常 ...... ");
        }

    }
}
